import torch as pt

pt.manual_seed(777)

x = pt.arange(0, 170*20).view(170, 20)


class MyModel(pt.nn.Module):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.em = pt.nn.Embedding(170*20, 300)
        self.rnn = pt.nn.RNN(300, 128, 2, batch_first=True)
        self.lin = pt.nn.Linear(128, 10)

    def forward(self, input):
        x = input
        print('input', x.shape)
        x = self.em(x)
        print('after embedding', x.shape)
        x, _ = self.rnn(x)
        print('after rnn', x.shape)
        x = x.reshape(-1, 128)
        print('before lin', x.shape)
        x = self.lin(x)
        print('after lin', x.shape)
        return x


model = MyModel()
model(x)
